# __init__.py
import os
import sys
from pathlib import Path

# 使用传统方式，但更加健壮
PROJECT_ROOT_STR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
if PROJECT_ROOT_STR not in sys.path:
    sys.path.insert(0, PROJECT_ROOT_STR)
    
from config.path_config import *
"""PPQ量化数据集模块"""
import torch
import torch.utils.data
from torchvision import transforms
from typing import Optional
from dataloader.dataloader import ImageFolderNoClass
from config.quantization_config import QuantizationConfig

class CalibrationDatasetBuilder:
    """校准数据集构建器"""
    
    def __init__(self, config: QuantizationConfig):
        self.config = config
        self.preprocess = transforms.Compose([
            transforms.Resize((640, 640)),
            transforms.ToTensor(),
        ])
    
    def build_dataset(self) -> torch.utils.data.DataLoader:
        """构建校准数据集"""
        if not os.path.isdir(self.config.data_dir):
            raise FileNotFoundError(f'数据目录未找到: {self.config.data_dir}')
        
        try:
            calib_dataset = ImageFolderNoClass(
                self.config.data_dir, 
                transform=self.preprocess
            )
            
            calib_loader = torch.utils.data.DataLoader(
                calib_dataset, 
                batch_size=self.config.batch_size, 
                pin_memory=True, 
                num_workers=self.config.num_workers
            )
            
            return calib_loader
            
        except Exception as e:
            raise RuntimeError("校准数据集加载失败: " + str(e))
    
    def collate_fn(self, batch: torch.Tensor) -> torch.Tensor:
        """数据整理函数"""
        return batch.to(self.config.device)